Generator module

@synalinks_export(["synalinks.modules.Generator", "synalinks.Generator"])
class Generator(Module):
    """
    Use a `LanguageModel` (or a `DecisionModel`) to generate a data model from an
    arbitrary input data model.

    Example:

    ```python
    import synalinks
    import asyncio

    async def main():

        class Query(DataModel):
            query: str = synalinks.Field(
                description="The user query",
            )

        class AnswerWithCritique(synalinks.DataModel):
            thinking: str = synalinks.Field(
                description="Your step by step thinking",
            )
            critique: str = synalinks.Field(
                description="The critique of the above thinking",
            )
            answer: str = synalinks.Field(
                description="The correct answer",
            )

        language_model = synalinks.LanguageModel(
            model="ollama/mistral",
        )

        x0 = synalinks.Input(data_model=Query)
        x1 = await synalinks.Generator(
            data_model=AnswerWithCritique,
            language_model=language_model,
        )(x0)

        program = synalinks.Program(
            inputs=x0,
            outputs=x1,
            name="chain_of_thought_with_critique",
            description="Useful to answer step by step and evaluate your answer",
        )

    if __name__ == "__main__":
        asyncio.run(main())
    ```

    Pass a `decision_model` to answer with a `DecisionModel` instead of the
    language model: it is faster and cheaper, and its answers are calibrated,
    but it does not generate. Every field of the data model must be a question
    it can answer (a `bool`, a string enum, a score such as `synalinks.Rating`,
    or a `score_schema` object, see `DecisionModel`), otherwise an
    `UnsupportedSchemaError` is raised. The instructions and examples are sent
    to it as the system message.

    ```python
    import synalinks
    import asyncio
    from typing import Literal

    async def main():

        class Ticket(synalinks.DataModel):
            message: str = synalinks.Field(
                description="The customer message",
            )

        class Triage(synalinks.DataModel):
            is_billing: bool = synalinks.Field(
                description="Is the ticket about billing?",
            )
            urgency: Literal["low", "medium", "high"] = synalinks.Field(
                description="How urgent is the ticket?",
            )

        decision_model = synalinks.DecisionModel(
            model="typesafe/jev-latest",
        )

        x0 = synalinks.Input(data_model=Ticket)
        x1 = await synalinks.Generator(
            data_model=Triage,
            decision_model=decision_model,
            instructions="Triage the support tickets of an online shop.",
        )(x0)

        program = synalinks.Program(
            inputs=x0,
            outputs=x1,
            name="ticket_triage",
            description="Triage the support tickets",
        )

    if __name__ == "__main__":
        asyncio.run(main())
    ```

    Args:
        schema (dict): The target JSON schema.
            If not provided use the `data_model` to infer it.
        data_model (DataModel | SymbolicDataModel | JsonDataModel): The target data
            model for structured output.
        language_model (LanguageModel): The language model to use.
        prompt_template (str): The jinja2 prompt template.
        prompt_variables (dict): Optional. Extra static Jinja2 variables
            rendered into the prompt template on every call, alongside the
            built-ins (``instructions``, ``examples``, ``inputs_schema``,
            ``outputs_schema``, which win on a name clash). Plain
            configuration, not trainable state, so in-context optimization
            never rewrites them; agents use this channel for AGENTS.md
            conventions (``custom_instructions``), the Agent Skills listing
            (``available_skills``) and RLM's sandbox functions
            (``sandbox_functions``).
        examples (list): The default list of examples, the examples
            are a list of tuples containing input/output JSON pairs.
        instructions (str): The default instructions being a string containing
            instructions for the language model.
        seed_instructions (list): Optional. A list of instructions to use as seed for the
            optimization. If not provided, use the default instructions as seed.
        use_inputs_schema (bool): Optional. Whether or not use the inputs schema in
            the prompt (Default to False).
        use_outputs_schema (bool): Optional. Whether or not use the outputs schema in
            the prompt (Default to False).
        return_inputs (bool): Optional. Whether or not to concatenate the inputs to
            the outputs (Default to False).
        temperature (float): Optional. The sampling temperature for the LM call.
            Default to None: when None it is NOT sent, so the model's own
            generation defaults apply (e.g. a vLLM-served model uses its
            `generation_config.json`). Set a float to override.
        max_tokens (int): Optional. Cap on the number of tokens generated. Default
            to None (not sent → the provider/model default). Set it to bound
            runaway / looping generations.
        top_p (float): Optional. Nucleus-sampling top-p. Default None (not sent →
            model default).
        top_k (int): Optional. Top-k sampling. Default None (not sent → model
            default).
        reasoning_effort (string): Optional. The reasoning effort for the LM call
            between ['minimal', 'low', 'medium', 'high', 'disable', 'none', None].
            Default to None (no reasoning).
        streaming (str): Optional. If true stream the LM response, enabled only if
            `schema` is `None`. Honored in every phase (inference, reward,
            optimizer) and in training: in a batched loop (predict / evaluate /
            the optimizer forward pass) or during training the stream is drained
            into a concrete prediction so it stays scorable while still recording
            time-to-first / time-to-last token (the optimizer phase included, so
            its TTFT is measured). Only an interactive single call (no active
            op_scope and not training) hands the live stream back to the caller.
        tools (list): Optional. Live `synalinks.modules.Tool` objects (or a
            `{name: Tool}` mapping) the generator always exposes, merged with any
            `tools` passed to `call`. Serialized as `tool_schemas` (their wire
            form) in the config, the way `data_model` is stored as `schema`.
        tool_schemas (list): Optional. Already-wire-formatted tool declaration
            dicts (OpenAI `{"type": "function", ...}` shape) the generator always
            exposes. Merged with any `tool_schemas` passed to `call`. Being plain
            JSON, they serialize with the config (unlike per-call `tools`).
        name (str): Optional. The name of the module.
        description (str): Optional. The description of the module.
        trainable (bool): Whether the module's variables should be trainable.
        decision_model (DecisionModel): Optional. A decision model to answer
            with instead of the language model. Every field of the schema must
            be a question it can answer (see `DecisionModel.check_schema()`).
            Without `decision_model` nor `language_model`, the default decision
            model (see `synalinks.set_default_decision_model`) is used when it
            can answer the schema, else the default language model.
    """

    def __init__(
        self,
        *,
        schema=None,
        data_model=None,
        language_model=None,
        prompt_template=None,
        prompt_variables=None,
        examples=None,
        instructions=None,
        seed_instructions=None,
        use_inputs_schema=False,
        use_outputs_schema=False,
        return_inputs=False,
        temperature=None,
        max_tokens=None,
        top_p=None,
        top_k=None,
        reasoning_effort=None,
        streaming=False,
        tools=None,
        tool_schemas=None,
        name=None,
        description=None,
        trainable=True,
        decision_model=None,
    ):
        super().__init__(
            name=name,
            description=description,
            trainable=trainable,
        )
        if not schema and data_model:
            schema = data_model.get_schema()
        self.schema = schema
        self.language_model = _get_lm(language_model)
        self.decision_model = resolve_decision_model(decision_model, language_model)
        if self.decision_model is not None:
            try:
                self.decision_model.check_schema(self.schema)
            except UnsupportedSchemaError:
                if decision_model is not None:
                    raise
                # The default decision model cannot answer this schema (it
                # needs generation): use the language model.
                self.decision_model = None
        if not prompt_template:
            prompt_template = default_prompt_template()
        self.prompt_template = prompt_template
        self.prompt_variables = dict(prompt_variables or {})
        if not examples:
            examples = []
        self.examples = examples
        if not instructions and self.schema:
            data_model_keys = list(self.schema["properties"].keys())
            instructions = default_instructions(data_model_keys)
        self.instructions = instructions
        self.return_inputs = return_inputs
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.top_p = top_p
        self.top_k = top_k
        efforts = ["minimal", "low", "medium", "high", "xhigh", "disable", "none", None]
        if reasoning_effort not in efforts:
            raise ValueError(
                f"The reasoning effort parameter should be one of: {efforts}"
            )
        self.reasoning_effort = reasoning_effort
        self.use_inputs_schema = use_inputs_schema
        self.use_outputs_schema = use_outputs_schema
        if schema and streaming:
            streaming = False
        self.streaming = streaming
        # Live `Tool` objects the generator always exposes; merged with any
        # passed per-call. They are not JSON, so `get_config` serializes them as
        # wire-format `tool_schemas` (the way `data_model` is stored as `schema`).
        self.tools = _as_tool_list(tools)
        # Already wire-format (OpenAI `{"type": "function", ...}`) tool
        # declarations the generator always exposes. Plain JSON dicts, so they
        # serialize as-is and are merged with any per-call `tool_schemas`.
        self.tool_schemas = tool_schemas

        predictions = [
            Prediction(
                inputs=example[0],
                outputs=example[1],
                reward=None,
            ).get_json()
            for example in examples
        ]

        if not seed_instructions:
            seed_instructions = []
        self.seed_instructions = seed_instructions

        seed_candidates = [
            {
                "instructions": seed_instruction,
            }
            for seed_instruction in self.seed_instructions
        ]

        self.state = self.add_variable(
            initializer=Instructions(
                instructions=instructions,
                examples=predictions,
                seed_candidates=seed_candidates,
            ).get_json(),
            data_model=Instructions,
            name="state_" + self.name,
        )

    async def call(self, inputs, tools=None, tool_schemas=None, training=False):
        if not inputs:
            return None
        # Merge the always-on, constructor-level tools/schemas with any passed
        # per-call.
        tools = self.tools + _as_tool_list(tools) or None
        tool_schemas = list(self.tool_schemas or []) + list(tool_schemas or []) or None
        msgs = self.format_messages(inputs)
        # Streaming is honored in every phase (inference, reward, optimizer) and
        # in training. When it is drained below (any batched loop or training
        # pass) the prediction stays concrete and scorable while still recording
        # TTFT / TTLT, so there is no reason to disable it -- including during
        # training and single-sample (batch=1) runs. Keeping it on in the
        # optimizer phase is deliberate: its time-to-first-token must be measured
        # too.
        streaming = bool(self.streaming)
        # Only forward sampling params that were explicitly set, so an unset
        # (None) value falls through to the model's own generation defaults
        # (e.g. a vLLM-served model applies its generation_config.json) instead
        # of being sent as `null`. litellm.drop_params drops *unsupported* keys,
        # not None-valued ones, so the filtering must happen here.
        sampling = {}
        if self.temperature is not None:
            sampling["temperature"] = self.temperature
        if self.max_tokens is not None:
            sampling["max_tokens"] = self.max_tokens
        if self.top_p is not None:
            sampling["top_p"] = self.top_p
        if self.top_k is not None:
            sampling["top_k"] = self.top_k
        model = self.decision_model or self.language_model
        value = await model(
            msgs,
            schema=self.schema,
            tools=tools,
            tool_schemas=tool_schemas,
            streaming=streaming,
            reasoning_effort=self.reasoning_effort,
            **sampling,
        )
        if isinstance(value, StreamingIterator):
            # A purely interactive call (no op_scope, not training) hands the
            # lazy stream to the caller to consume. Inside a batched loop
            # (predict / evaluate run under op_scope) or during training the
            # prediction must be concrete -- to be scored and recorded into the
            # optimizer's prediction state -- so drain the stream into a
            # ChatMessage; time-to-first / time-to-last token are still recorded
            # as it is consumed.
            if current_op_scope() is None and not training:
                return value
            value = await value.aconsume(name=f"{self.language_model.name}_response")
        if not value:
            result = None
        else:
            result = value.clone(name="prediction_" + self.name)
        if result:
            if training:
                predictions = self.state.get("current_predictions")
                predictions.append(
                    {
                        "inputs": inputs.get_json(),
                        "outputs": result.get_json(),
                        "reward": None,
                    }
                )
            if self.return_inputs:
                return await ops.concat(
                    inputs,
                    result,
                    name="with_inputs_" + self.name,
                )
            else:
                return result
        return None

    async def compute_output_spec(
        self, inputs, tools=None, tool_schemas=None, training=False
    ):
        if self.schema:
            if self.return_inputs:
                return await ops.concat(
                    inputs,
                    SymbolicDataModel(
                        schema=self.schema,
                        name=self.name,
                    ),
                    name="with_inputs_" + self.name,
                )
            else:
                return SymbolicDataModel(
                    schema=self.schema,
                    name=self.name,
                )
        else:
            if self.return_inputs:
                return await ops.concat(
                    inputs,
                    SymbolicDataModel(
                        schema=ChatMessage.get_schema(),
                        name=self.name,
                    ),
                    name="with_inputs_" + self.name,
                )
            else:
                return SymbolicDataModel(
                    schema=ChatMessage.get_schema(),
                    name=self.name,
                )

    def format_messages(self, inputs=None):
        return format_chat_messages(inputs, self._render_system_message(inputs))

    def _render_system_message(self, inputs):
        template = jinja2.Template(self.prompt_template)
        rendered_prompt = template.render(
            **{
                # Built-in variables win over a same-named prompt variable.
                **self.prompt_variables,
                "inputs_schema": (
                    inputs.get_schema() if self.use_inputs_schema else None
                ),
                "outputs_schema": self.schema if self.use_outputs_schema else None,
                "examples": [
                    (pred.get("inputs"), pred.get("outputs"))
                    for pred in self.state.get("examples")
                ],
                "instructions": self.state.get("instructions"),
            }
        )
        return ChatMessage(role="system", content=rendered_prompt)

    def get_config(self):
        config = {
            "schema": self.schema,
            "prompt_template": self.prompt_template,
            "prompt_variables": self.prompt_variables,
            "examples": self.examples,
            "instructions": self.instructions,
            "seed_instructions": self.seed_instructions,
            "use_inputs_schema": self.use_inputs_schema,
            "use_outputs_schema": self.use_outputs_schema,
            "return_inputs": self.return_inputs,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "reasoning_effort": self.reasoning_effort,
            "streaming": self.streaming,
            # Live `tools` are converted to their wire form and stored alongside
            # `tool_schemas` (the way `data_model` is stored as `schema`).
            "tool_schemas": (
                list(self.tool_schemas or []) + _tools_to_schemas(self.tools)
            )
            or None,
            "name": self.name,
            "description": self.description,
            "trainable": self.trainable,
        }
        language_model_config = {
            "language_model": serialization_lib.serialize_synalinks_object(
                self.language_model,
            )
        }
        if self.decision_model is not None:
            language_model_config["decision_model"] = (
                serialization_lib.serialize_synalinks_object(self.decision_model)
            )
        return {
            **config,
            **language_model_config,
        }

    @classmethod
    def from_config(cls, config):
        language_model = serialization_lib.deserialize_synalinks_object(
            config.pop("language_model"),
        )
        if "decision_model" in config:
            config["decision_model"] = serialization_lib.deserialize_synalinks_object(
                config.pop("decision_model"),
            )
        return cls(
            language_model=language_model,
            **config,
        )